package com.ddone.demo;

import javolution.text.Text;
import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils;

import java.util.ArrayList;
import java.util.List;

/**
 * @author ddone
 * @date 2024/10/15-23:56
 * 接收一组 AOI（多边形）坐标，向外扩展一定的距离，并返回新的扩展后的 AOI 坐标。
 *  1. 接收一组 AOI 的坐标：AOI 通常由一组点（经纬度）定义。
 *  2. 每个点向外扩展：根据输入的距离，将每个点都向外扩展指定的距离。
 *  3. 返回新的扩展后的 AOI：以新的坐标返回这个多边形。
 */
@Description(name = "expand_aoi", value = "_FUNC_(Array<String> coordinates, double distanceInKm) - returns expanded AOI coordinates")
public class ExpandAoiHiveFunction extends GenericUDF {
    private static final double EARTH_RADIUS = 6371.009;
    // 更加精确的地球半径，单位：公里
    private PrimitiveObjectInspector distanceOI;
    private PrimitiveObjectInspector stringAOI;

    // 将角度转换为弧度
    private static double rad(double d) {
        return d * Math.PI / 180.0;
    }

    // 将弧度转换为角度
    private static double deg(double rad) {
        return rad * 180.0 / Math.PI;
    }

    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        if (arguments.length != 2) {
            throw new UDFArgumentLengthException("The function expand_aoi requires 2 arguments.");
        }

        // 检查输入参数类型
        if (!(arguments[0] instanceof PrimitiveObjectInspector)) {
            throw new UDFArgumentTypeException(0, "The first argument must be an array of strings.");
        }

        if (!(arguments[1] instanceof PrimitiveObjectInspector)) {
            throw new UDFArgumentTypeException(1, "The second argument must be a double representing distance in kilometers.");
        }

        stringAOI = (PrimitiveObjectInspector) arguments[0];
        distanceOI = (PrimitiveObjectInspector) arguments[1];

        // 返回数组，数组中的元素是经纬度的字符串格式
        return ObjectInspectorFactory.getStandardListObjectInspector(PrimitiveObjectInspectorFactory.writableStringObjectInspector);
    }

    @Override
    public Object evaluate(DeferredObject[] arguments) throws HiveException {
        // 读取AOI坐标数组
        List<Text> coordinates = (List<Text>) arguments[0].get();
        double distanceInKm = PrimitiveObjectInspectorUtils.getDouble(arguments[1].get(), distanceOI);

        List<Text> expandedAOI = new ArrayList<>();
        double angularDistance = distanceInKm / EARTH_RADIUS;

        for (Text coordinate : coordinates) {
            String[] parts = coordinate.toString().split("_");
            if (parts.length != 2) {
                throw new HiveException("Each coordinate must be in the format 'lon_lat'.");
            }

            double lon = Double.parseDouble(parts[0]);
            double lat = Double.parseDouble(parts[1]);

            double latRad = rad(lat);
            double lonRad = rad(lon);

            // 在每个方向扩展距离，简单实现为一个方向上的扩展
            double latNew = Math.asin(Math.sin(latRad) * Math.cos(angularDistance)
                    + Math.cos(latRad) * Math.sin(angularDistance) * Math.cos(0));

            double lonNew = lonRad + Math.atan2(Math.sin(0) * Math.sin(angularDistance) * Math.cos(latRad),
                    Math.cos(angularDistance) - Math.sin(latRad) * Math.sin(latNew));

            expandedAOI.add(new Text(deg(lonNew) + "_" + deg(latNew)));
        }

        return expandedAOI;
    }

    @Override
    public String getDisplayString(String[] children) {
        return "expand_aoi(" + children[0] + ", " + children[1] + ")";
    }
}
